// 
// task.c - Implements the functionality needed to multitask.
//          Written for JamesM's kernel development tutorials.
//

#include "task.h"
#include "paging.h"
#include "util.h"
#include "kheap.h"
#include "tables.h"

// The currently running task.
volatile task_t *current_task;

// The start of the task linked list.
volatile task_t *ready_queue;

// Some externs are needed to access members in paging.c...
extern page_directory_t *kernel_directory;
extern page_directory_t *current_directory;
extern void alloc_frame(page_t*,int,int);
extern unsigned int initial_esp;
extern unsigned int read_eip();

// The next available process ID.
unsigned int next_pid = 1;

void initialize_tasking()
{
	// Rather important stuff happening, no interrupts please!
	disable();

	// Relocate the stack so we know where it is.
	move_stack((void*)0xE0000000, 0x2000);

	// Initialise the first task (kernel task)
	current_task = ready_queue = (task_t*)kmalloc(sizeof(task_t));
	current_task->id = next_pid++;
	current_task->esp = current_task->ebp = 0;
	current_task->eip = 0;
	current_task->page_directory = current_directory;
	current_task->next = 0;
	current_task->kernel_stack = kmalloc_a(KERNEL_STACK_SIZE);
	
	// Reenable interrupts.
	enable();
}








int paging_createDirectory( struct task_t *p )
{
	int index;
	struct page_directory_t * pd;
	
	// lock critical section
	//mutex_lock( &paging_lock );
	
	// allocate some physical memory for the page directory
	p->page_dir = (struct PAGE_DIRECTORY *)physical_pageAlloc();
	if( p->page_dir == NULL )
	{
		mutex_unlock( &paging_lock );
		return FAIL;
	}
	// map the p->page_dir physical address into the current address space so we can read/write to it
	pd = (struct PAGE_DIRECTORY *)paging_mapQuick( p->page_dir );
	// clear out the page directory
	memset( pd, 0x00, sizeof(struct PAGE_DIRECTORY) );
	// set some default entrys in the page directory
	for( index=0 ; index<PAGE_ENTRYS; index++ )
	{
		// get the next entry
		struct PAGE_DIRECTORY_ENTRY * pde = &pd->entry[index];
		// set the privilege ala the process
		pde->user = p->privilege;
		// if its the last entry
		if( index == PAGE_ENTRYS-1 )
		{
			// store the page dir as the last entry in itself (fractal mapping)
			pde->present	= TRUE;
			pde->readwrite	= READWRITE;
			pde->address	= TABLE_SHIFT_R( p->page_dir );		
		}
		else
		{
			pde->present   = FALSE;
			pde->readwrite = READONLY;
		}
	}
	// unlock critical section
	mutex_unlock( &paging_lock );
	// return success
	return SUCCESS;
}





























int new_task(void (*entry)(void))
{
	disable();

	// Clone the address space.
	page_directory_t *directory = clone_directory(current_directory);

	// Create a new process.
	task_t *new_task = (task_t*)kmalloc(sizeof(task_t));

	new_task->id = next_pid++;
	new_task->esp = new_task->ebp = 0;
	new_task->eip = 0;
	new_task->page_directory = directory;
	current_task->kernel_stack = kmalloc_a(KERNEL_STACK_SIZE);
	new_task->next = 0;

	// Add it to the end of the ready queue.
	task_t *tmp_task = (task_t*)ready_queue;
	while (tmp_task->next)
		tmp_task = tmp_task->next;
	tmp_task->next = new_task;

	unsigned int esp; asm volatile("mov %%esp, %0" : "=r"(esp));
	unsigned int ebp; asm volatile("mov %%ebp, %0" : "=r"(ebp));
	new_task->esp = esp;
	new_task->ebp = ebp;
	new_task->eip = entry;
	enable();

	return new_task->id;
}

void move_stack(void *new_stack_start, unsigned int size)
{
	unsigned int i;
	// Allocate some space for the new stack.
	for( i = (unsigned int)new_stack_start;
		i >= ((unsigned int)new_stack_start-size);
		i -= 0x1000)
	{
		// General-purpose stack is in user-mode.
		alloc_frame( get_page(i, 1, current_directory), 0 /* User mode */, 1 /* Is writable */ );
	}

	// Flush the TLB by reading and writing the page directory address again.
	unsigned int pd_addr;
	asm volatile("mov %%cr3, %0" : "=r" (pd_addr));
	asm volatile("mov %0, %%cr3" : : "r" (pd_addr));

	// Old ESP and EBP, read from registers.
	unsigned int stack; asm volatile("mov %%esp, %0" : "=r" (stack));
	unsigned int base;  asm volatile("mov %%ebp, %0" : "=r" (base));

	// Offset to add to old stack addresses to get a new stack address.
	unsigned int offset = (unsigned int)new_stack_start - initial_esp;

	// New ESP and EBP.
	unsigned int nstack = stack + offset;
	unsigned int nbase  = base  + offset;

	// Copy the stack.
	memcpy((void*)nstack, (void*)stack, initial_esp-stack);

	// Backtrace through the original stack, copying new values into
	// the new stack.  
	for(i = (unsigned int)new_stack_start; i > (unsigned int)new_stack_start-size; i -= 4)
	{
		unsigned int tmp = * (unsigned int*)i;
		// If the value of tmp is inside the range of the old stack, assume it is a base pointer
		// and remap it. This will unfortunately remap ANY value in this range, whether they are
		// base pointers or not.
		if (( stack < tmp) && (tmp < initial_esp))
		{
			tmp = tmp + offset;
			unsigned int *tmp2 = (unsigned int*)i;
			*tmp2 = tmp;
		}
	}

	// Change stacks.
	asm volatile("mov %0, %%esp" : : "r" (nstack));
	asm volatile("mov %0, %%ebp" : : "r" (nbase));
}

void switch_task()
{
	// If we haven't initialised tasking yet, just return.
	if (!current_task)
	return;

	// Read esp, ebp now for saving later on.
	unsigned int esp, ebp, eip;
	asm volatile("mov %%esp, %0" : "=r"(esp));
	asm volatile("mov %%ebp, %0" : "=r"(ebp));

	// Read the instruction pointer. We do some cunning logic here:
	// One of two things could have happened when this function exits - 
	//   (a) We called the function and it returned the EIP as requested.
	//   (b) We have just switched tasks, and because the saved EIP is essentially
	//       the instruction after read_eip(), it will seem as if read_eip has just
	//       returned.
	// In the second case we need to return immediately. To detect it we put a dummy
	// value in EAX further down at the end of this function. As C returns values in EAX,
	// it will look like the return value is this dummy value! (0x12345).
	eip = read_eip();

	// Have we just switched tasks?
	if (eip == 0x12345)
		return;

	// No, we didn't switch tasks. Let's save some register values and switch.
	current_task->eip = eip;
	current_task->esp = esp;
	current_task->ebp = ebp;

	// Get the next task to run.
	current_task = current_task->next;
	// If we fell off the end of the linked list start again at the beginning.
	if (!current_task) current_task = ready_queue;

	eip = current_task->eip;
	esp = current_task->esp;
	ebp = current_task->ebp;

	// Make sure the memory manager knows we've changed page directory.
	current_directory = current_task->page_directory;
	
	set_kernel_stack(current_task->kernel_stack + KERNEL_STACK_SIZE);
	
	// Here we:
	// * Stop interrupts so we don't get interrupted.
	// * Temporarily puts the new EIP location in ECX.
	// * Loads the stack and base pointers from the new task struct.
	// * Changes page directory to the physical address (physicalAddr) of the new directory.
	// * Puts a dummy value (0x12345) in EAX so that above we can recognise that we've just
	//   switched task.
	// * Restarts interrupts. The STI instruction has a delay - it doesn't take effect until after
	//   the next instruction.
	// * Jumps to the location in ECX (remember we put the new EIP in there).
	asm volatile("         \
	cli;                 \
	mov %0, %%ecx;       \
	mov %1, %%esp;       \
	mov %2, %%ebp;       \
	mov %3, %%cr3;       \
	mov $0x12345, %%eax; \
	sti;                 \
	jmp *%%ecx           "
		 : : "r"(eip), "r"(esp), "r"(ebp), "r"(current_directory->physicalAddr));
}

int fork()
{
	disable();

	// Take a pointer to this process' task struct for later reference.
	task_t *parent_task = (task_t*)current_task;

	// Clone the address space.
	page_directory_t *directory = clone_directory(current_directory);

	// Create a new process.
	task_t *new_task = (task_t*)kmalloc(sizeof(task_t));

	new_task->id = next_pid++;
	new_task->esp = new_task->ebp = 0;
	new_task->eip = 0;
	new_task->page_directory = directory;
	current_task->kernel_stack = kmalloc_a(KERNEL_STACK_SIZE);
	new_task->next = 0;

	// Add it to the end of the ready queue.
	task_t *tmp_task = (task_t*)ready_queue;
	while (tmp_task->next)
		tmp_task = tmp_task->next;
	tmp_task->next = new_task;

	// entry point for the new process
	unsigned int eip = read_eip();

	// We could be the parent or the child here - check.
	if (current_task == parent_task)
	{
		// We are the parent, so set up the esp/ebp/eip for our child.
		unsigned int esp; asm volatile("mov %%esp, %0" : "=r"(esp));
		unsigned int ebp; asm volatile("mov %%ebp, %0" : "=r"(ebp));
		new_task->esp = esp;
		new_task->ebp = ebp;
		new_task->eip = eip;
		enable();

		return new_task->id;
	}
	else
	{
		// We are the child.
		return 0;
	}
}

int getpid()
{
	return current_task->id;
}

void user_mode()
{
	// Set up our kernel stack.
	set_kernel_stack(current_task->kernel_stack+KERNEL_STACK_SIZE);
	
	// Set up a stack structure for switching to user mode.
	asm volatile("  \
		cli; \
		mov $0x23, %ax; \
		mov %ax, %ds; \
		mov %ax, %es; \
		mov %ax, %fs; \
		mov %ax, %gs; \
			    \
		\
		mov %esp, %eax; \
		pushl $0x23; \
		pushl %esp; \
		pushf; \
		pop %ebx; \
		or $0x200, %ebx; \
		push %ebx; \
		pushl $0x1B; \
		push $1f; \
		iret; \
		1: \
	");
}
